fix: regenerate datasets left on disk by a small-datasets run - #471

Fixes the root cause behind PyAutoLabs/autolens_workspace_test#260.

The bug

should_simulate was existence-only and asymmetric. It force-regenerated when PYAUTO_SMALL_DATASETS=1, but had no corresponding check on the full-resolution path:

if os.environ.get("PYAUTO_SMALL_DATASETS") == "1":
    if Path(dataset_path).exists():
        shutil.rmtree(dataset_path)

return not Path(dataset_path).exists()

dataset/ is gitignored in the workspaces, so CI clones fresh, always simulates, and can never hit this. Locally the directory persists indefinitely — and since PYAUTO_SMALL_DATASETS=1 is the default for most harness runs, a single earlier run leaves 16×16 FITS on disk that every later full-resolution run then loads silently.

The result was deterministic, bit-identical assertion failures in three jax_grad scripts that pass in CI on the same commit. That is an active trap rather than a cosmetic one: the failures look exactly like fresh correctness regressions (pure_callback constant-folding, an AD-vs-FD tolerance breach), and were initially reported as such.

The fix

The regime is recorded nowhere on disk, so it is inferred from data.fits: the cap in Mask2D.circular / Grid2D.uniform rewrites anything larger to exactly SMALL_DATASETS_SHAPE_NATIVE, so a data.fits at exactly (16, 16) can only have come from a capped run.

The predicate ends in shutil.rmtree, so it is deliberately narrow — a false positive silently deletes real data:

• Exactly the cap shape, never "at or below" it. The cap cannot emit 12×12, so widening the test buys no detection and only adds risk.
• data.fits by name, never "the first FITS in the directory". PSF kernels are legitimately tiny at full resolution — 11×11 appears in 49 places in autolens_workspace, and autolens_workspace_test/dataset/cluster/test/psf.fits is on disk at 5760 B, byte-identical in size to a capped data.fits. A glob-based check would regenerate every dataset carrying a PSF on every single run.
• Unknown means no. A missing, unreadable or non-2D data.fits returns False, preserving existing existence-only behaviour.

Scope — imaging only

This covers the imaging manifestation, not the whole bug class. Two families stay exposed and regress to existence-only, documented in the should_simulate docstring rather than left to be rediscovered:

• point-source and weak-lensing datasets are JSON with no FITS;
• interferometer datasets keep their shape under the cap (visibility count fixed by the uv file, real-space grid capped behind it), so a capped run writes identical NAXIS with different values — and therefore fails silently, with no assertion to trip.

Closing those needs the regime recorded at write time rather than inferred at read time: PyAutoLabs/PyAutoNerves#153. Kept out of this PR deliberately — it changes a header card on every FITS the stack writes, which could disturb round-trip tests and file-hash pins.

A separate live defect found on the way is filed as #470: the small-datasets branch rmtrees dataset/point_source/simple, which is committed and explicitly allowlisted in that repo's .gitignore.

Cost

One FITS header read, 0.64 ms warm, and only when a data.fits exists. No re-simulation in the steady state. (The first call in a process also pays a ~206 ms astropy.io.fits import that import autolens does not itself trigger; any script reaching this point loads FITS moments later regardless, so it is paid earlier rather than added.)

Unconditional regeneration was considered and rejected on cost: a full simulate is ~8.7 s, 17 of 24 curated smoke entries in autolens_workspace_test are full-regime and many share dataset directories, and end users of autolens_workspace are on the full path by default across ~247 call sites — so every script would re-simulate on every run.

Testing

• 1168 tests pass (full test_autoarray suite).
• 12 new tests: all four regime transitions — the bug was precisely that one of the four was never exercised — plus five false-positive guards, each asserting a dataset is preserved (tiny PSF, PSF at exactly the cap shape, 12×12 data, JSON-only dataset, unreadable data.fits).
• Control-tested: the small→full regression test returns False against the unfixed body and True against the fixed one.
• Verified end to end against the original failure: poison the dataset with one PYAUTO_SMALL_DATASETS=1 run (data.fits → 5760 B), then run imaging/jax_grad/lp.py under the full_datasets profile. Before: AssertionError: All source-parameter gradients are ~zero, source gradients at ~1e-12. After: regenerated at 262080 B, checks pass.

Downstream

No public API change. No workspace change is required, and no assertion tolerance was changed — the three jax_grad assertions were correct throughout and were detecting a genuinely invalid dataset.

🤖 Generated with Claude Code

https://claude.ai/code/session_01VEHLT33XpVcRt5YCJGLRMJ

Generated by Claude Code